
from Bio.Align import substitution_matrices
from align import Table,SubstMatrix

# aligner = PairwiseAligner()

# alignments = aligner.align("ATAGAGAATAG","ATGGCAGATAGA" )

# aligner
# alignments = aligner.align("ANDDR","AARRD") 

# print(len(alignments))

def create_mat(nrows, ncols):
    mat = []
    for i in range(nrows):
        mat.append([])
        for j in range(ncols):
            mat[i].append(0)
    return mat

def dotplot(seq1, seq2):
    mat = create_mat(len(seq1), len(seq2))
    for i in range( len(seq1)):
        for j in range( len(seq2)):
            if seq1[i] == seq2[j]:
                mat[i][j] = 1 
    return mat

def print_dotplot(mat, s1, s2):
    import sys
    sys.stdout.write(" "+s2+"\n")
    for i in range( len(mat)):
        sys.stdout.write(s1[i])
        for j in range( len(mat[i])):
            if mat[i][j] >= 1:
                sys.stdout.write("*" )
            else:
                sys.stdout.write(" ")
        sys.stdout.write("\n")

def extended_dotplot (seq1, seq2, window, stringency):
    mat = create_mat(len(seq1), len(seq2))
    start = int(window/2)
    for i in range(start, len(seq1)-start):
        for j in range(start, len(seq2)-start):
            matches = 0
            l=j- start 
            for k in range(i-start, i+start+1):
                if seq1[k] == seq2[l]:
                    matches += 1
                    l+=1
                    if matches >= stringency: 
                        mat[i][j] = 1 
    return mat

def create_submat(match, mismatch, alphabet):
    sm = {}
    for c1 in alphabet:
        for c2 in alphabet:
            if (c1 == c2):
                sm[c1+c2] = match
            else:
                sm[c1+c2] = mismatch
    return sm

def read_submat_file(filename):
    '''
    read tsv file, split by \\t
    first row is alphabet, example:

    A   B   C

    1   2   3

    2   3   4

    3   4   5
    '''
    
    sm = {} 
    with open(filename, "r") as f:
        line = f.readline().strip()
        alphabet  = line.split("\t")
        for from_char in alphabet :
            line = f.readline().strip().split("\t")
            for i,to_char in enumerate(alphabet) :
                sm[from_char+to_char] = line[i]
    return sm

def indentity(seq1,seq2,alphabet="ACGT"):
    sm = SubstMatrix().create_submat(1,0,alphabet)
    table = Table(seq1,seq2)
    table.needleman_Wunsch(sm,0)
    equal = table[-1,-1].value
    return equal/max(len(seq1),len(seq2))


if __name__ == "__main__":
    s1 = "CGATATAG"
    s2 = "TATATATT"
    mat1 = dotplot(s1, s2)
    print_dotplot(mat1, s1, s2)

    s1 = "CGATATAGATT"
    s2 = "TATATAGTAT"
    mat2 = extended_dotplot(s1, s2, 5, 4)
    print_dotplot(mat2, s1, s2)

    print(create_submat(1,0,"ACGT"))
    substitute_mat = substitution_matrices.load("BLOSUM62")
    print(substitute_mat.alphabet)
    print(substitute_mat.tolist())
    print(substitute_mat[2,2])

    # print(read_submat_file("BLOSUM62.txt"))
    print(indentity("ATAT","ATAT"))